# 脚本从外接收参数
# python3 main.py aa.csv bb.csv

import csv
import sqlite3
import os
import sys
from itertools import islice

db_file = 'vuln.db'

csv_file = sys.argv[1]
new_file = sys.argv[2]

db = sqlite3.connect(db_file)
cursor = db.cursor()


def select(plugin_id):
    sql = 'SELECT Plugin_ID, NAME, Risk, Description, Solution FROM VULNDB where Plugin_ID="{}"'.format(plugin_id)
    result = cursor.execute(sql).fetchone()
    return result


def join_list(*args):
    """
    传递多个值，一次性把值全部变为列表
    :param args:
    :return: 列表
    """
    tmp_list = list()
    for var in args:
        tmp_list.append(var)
    return tmp_list


def read_file(file):

    reader = csv.reader(file)
    # old_title = ['plugin_id', 'cve', 'cvss', 'risk', 'host', 'protocol', 'port', 'name', 'synopsis', 'description','solution', 'see_also', 'plugin_output']

    info_list = list()

    for row in islice(reader, 1, None):
        plugin_id = row[0]
        risk = row[3]
        host = row[4]
        protocol = row[5]
        port = row[6]
        name = row[7]
        description = row[9]
        solution = row[10]

        translate_data = select(plugin_id)
        if translate_data:
            name = translate_data[1]
            risk = translate_data[2]
            description = translate_data[3]
            solution = translate_data[4]

        info_list.append(join_list(host, port, protocol, plugin_id, name, risk, description, solution))
    return info_list


def main():

    if not os.path.isfile(csv_file):
        print('[{}]文件不存在'.format(csv_file))
        return

    if os.path.isfile(new_file):
        print('[{}]文件已存在'.format(new_file))
        return

    with open(csv_file, mode='r', encoding='utf-8') as t1, open(new_file, mode='a', encoding='utf-8', newline='') as t2:

        # 目标文件title
        new_title = ['ip', '端口', '协议', '漏洞代码', '漏洞名称', '漏洞等级', '描述信息', '解决办法']

        # 写入头部
        t2_writer = csv.writer(t2)
        t2_writer.writerow(new_title)

        # 写入数据
        data_list = read_file(t1)
        print(data_list)
        t2_writer.writerows(data_list)


if __name__ == '__main__':
    try:
        main()
    finally:
        db.close()
